"""
分割分析服务

这是jinxiang_FF模块的主要接口，封装SegmentationAnalyzer，提供简化的文件夹分析功能。
"""

import os
import sys
import time
import numpy as np
from typing import Dict, Any, List, Optional
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.src import Classifier, get_optimal_device, prepare_output_directory


class ClassificationService:
    """
     图像分类
    """
    
    def __init__(self, model_path: str, device: Optional[str] = None):
        """
        初始化分割分析服务
        
        Args:
            model_path (str): 模型文件路径
            device (str, optional): 推理设备，如果为None则自动选择
        """
        self.model_path = model_path
        self.device = device or get_optimal_device()
        
        # 初始化分割分析器
        try:
            self.analyzer = Classifier(
                model_path=self.model_path,
                device=self.device
            )
            print(f"分割分析服务初始化成功，使用设备: {self.device}")
        except Exception as e:
            raise RuntimeError(f"初始化分割分析服务失败: {str(e)}")
    
    def analyze_oneimg(self, image_path: str, output_folder: str) -> Dict[str, Any]:
        """
        分析文件夹中的所有图像
        
        Args:
            input_folder (str): 输入文件夹路径
            output_folder (str): 输出文件夹路径
            target_classes (List[int]): 要分析的类别，默认为[1, 2]
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        result = {
            'success': False,
            'message': '',
            'result': 1.5,
           
        }
        classf=1.5
        try:

            # 分析单张图像（计时）
            start_time = time.time()
            analysis_result = self.analyzer.classify_single(image_path,output_folder)
            elapsed_analyze = time.time() - start_time
            print(f"  - 单图像分析耗时: {elapsed_analyze:.3f} 秒")


            '''
             analysis_result = {
            'class_id': predicted.item(),
            'class_name': self._get_class_name(predicted.item()),
            'confidence': confidence.item(),
            'probabilities': probabilities.cpu().numpy()[0]
        }
            '''

            outputimageStr = str(analysis_result['class_id']) + ',' + str(analysis_result['confidence'])
            # with open(fileoutput,mode='w',encoding = 'utf-8') as file_obj:
            #     file_obj.write(outputimageStr)
            if analysis_result['class_id'] == 0:
                classf = 0.5
            elif analysis_result['class_id'] == 1:
                classf = 1.0
            elif analysis_result['class_id'] == 2:
                classf = 1.5

            result = {
                'success': True,
                'message': '',
                'result': classf,
            
            }

            
            
        except Exception as e:
            print(f"处理图像失败 {image_path}: {str(e)}")
            result = {
            'success': False,
            'message': f"处理图像失败 {image_path}: {str(e)}",
            'result':1.5,
           
        }
             
        return result
    
    def _get_image_files(self, folder_path: str) -> List[str]:
        """
        获取文件夹中的所有图像文件
        
        Args:
            folder_path (str): 文件夹路径
            
        Returns:
            List[str]: 图像文件路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        image_files = []
        
        for file_path in Path(folder_path).iterdir():
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                image_files.append(str(file_path))
        
        return sorted(image_files)
    
    
    def _generate_summary_report(self, all_results: List[Dict], output_folder: str, target_classes: List[int]):
        """
        生成汇总报告到根目录的result.txt文件
        
        Args:
            all_results (List[Dict]): 所有分析结果
            output_folder (str): 输出文件夹
            target_classes (List[int]): 目标类别
        """
        try:
            summary_path = os.path.join(output_folder, 'result.txt')
            
            with open(summary_path, 'w', encoding='utf-8') as f:
                f.write("金相组织分析汇总报告\n")
                f.write("=" * 50 + "\n")
                f.write(f"分析时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"模型路径: {self.model_path}\n")
                f.write(f"目标类别: {target_classes}\n")
                f.write(f"总文件数: {len(all_results)}\n\n")
                
                # if all_results:
                #     # 获取类别名称映射（从第一个结果中获取）
                #     class_names = all_results[0].get('class_names', {})
                    
                #     # 计算平均占比
                #     avg_ratios = {}
                #     for class_id in target_classes:
                #         ratios = [r['class_ratios'].get(class_id, 0) for r in all_results]
                #         avg_ratios[class_id] = sum(ratios) / len(ratios)
                    
                #     f.write("整体统计:\n")
                #     for class_id, avg_ratio in avg_ratios.items():
                #         class_name = class_names.get(class_id, f"类别{class_id}")
                #         f.write(f"{class_name} (ID:{class_id}) 平均占比: {avg_ratio:.2f}%\n")
                #     f.write("\n")
                    
                #     f.write("详细结果:\n")
                #     for r in all_results:
                #         f.write(f"{r['file_name']} - ")
                #         for class_id in target_classes:
                #             ratio = r['class_ratios'].get(class_id, 0)
                #             class_name = class_names.get(class_id, f"类别{class_id}")
                #             f.write(f"{class_name}: {ratio:.2f}% ")
                #         f.write("\n")
            
            print(f"汇总报告已保存: {summary_path}")
            
        except Exception as e:
            print(f"生成汇总报告失败: {str(e)}")
    
    def _create_visualization(self, original_image_path: str, result: Dict[str, Any], output_path: str) -> bool:
        """
        创建简化的可视化结果，只保存overlap图像以提高性能
        
        Args:
            original_image_path: 原始图像路径
            result: 分析结果
            output_path: 输出路径
            
        Returns:
            bool: 是否保存成功
        """
        try:
            # 加载原始图像
            img = Image.open(original_image_path)
            
            # 获取结果数据
            class_ratios = result['class_ratios']
            class_masks = result['class_masks']
            
            # 创建overlap图像（P类和F类叠加）
            overlap_img = self._create_overlap_visualization(img, class_masks, class_ratios)
            
            if overlap_img is None:
                print(f"无法创建overlap图像: {output_path}")
                return False
            
            # 创建单个子图显示overlap图像
            fig, ax = plt.subplots(1, 1, figsize=(10, 8))
            
            # 显示overlap图像
            ax.imshow(overlap_img)
            ax.set_title('P+F Overlap')
            ax.axis('off')
            
            # 添加统计信息
            stats_text = f"P: {class_ratios.get(1, 0):.2f}%\nF: {class_ratios.get(2, 0):.2f}%"
            ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, 
                   fontsize=14, fontweight='bold', color='white',
                   bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.7),
                   verticalalignment='top')
            
            # 保存图像
            plt.tight_layout()
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
            plt.close()
            
            return True
            
        except Exception as e:
            print(f"创建可视化失败 {output_path}: {str(e)}")
            return False
    
    def _create_overlap_visualization(self, original_img: Image.Image, class_masks: Dict, class_ratios: Dict) -> Optional[np.ndarray]:
        """
        创建P类和F类的overlap可视化
        
        Args:
            original_img: 原始图像
            class_masks: 类别掩码字典
            class_ratios: 类别占比字典
            
        Returns:
            np.ndarray: overlap图像
        """
        try:
            # 将PIL图像转换为numpy数组
            if isinstance(original_img, Image.Image):
                img_array = np.array(original_img)
            else:
                img_array = original_img
            
            # 创建overlap图像
            overlap_img = img_array.copy()
            
            # 获取P类和F类的掩码
            if 1 not in class_masks or 2 not in class_masks:
                return None
                
            p_mask = class_masks[1]
            f_mask = class_masks[2]
            
            # 创建彩色overlay
            # P类用红色，F类用绿色，重叠区域用黄色
            p_overlay = np.zeros_like(img_array)
            f_overlay = np.zeros_like(img_array)
            
            # P类区域（红色）
            p_overlay[p_mask > 0] = [255, 0, 0]  # 红色
            # F类区域（绿色）
            f_overlay[f_mask > 0] = [0, 255, 0]  # 绿色
            
            # 重叠区域（黄色）
            overlap_mask = (p_mask > 0) & (f_mask > 0)
            p_overlay[overlap_mask] = [255, 255, 0]  # 黄色
            f_overlay[overlap_mask] = [255, 255, 0]  # 黄色
            
            # 将overlay与原图混合
            alpha = 0.6  # 透明度
            overlap_img = (1 - alpha) * img_array + alpha * (p_overlay + f_overlay)
            overlap_img = np.clip(overlap_img, 0, 255).astype(np.uint8)
            
            return overlap_img
            
        except Exception as e:
            print(f"Create overlap visualization failed: {str(e)}")
            return None